# -*- coding:utf-8 -*-
import os,sys
import re
import traceback
import time
import copy
from itertools import islice
from collections import deque
sys.path.append(os.path.join(os.path.abspath(os.path.dirname(__file__)), os.pardir, os.pardir))
import supeanut_config
sys.path.append(os.path.join(os.path.abspath(os.path.dirname(__file__)), os.pardir))
from CommonLib.mylog import mylog
from ExrightsTool import ExrightsTool
from BasicTool import BasicTool
from TecIndexRecog import TecIndexRecog
from StockTool import StockTool

'''
作者：supeanut
创建时间：2016-xx-xx xx:xx:xx
功能描述：
	形态匹配型 训练数据 生成
相关配置：
	supeanut_config.XXX
历史改动：
	2016-xx-xx: xxxxxx
'''

# 计算从date后一天`开盘买，持days天收盘卖的收益
# 使用ori行情数据，持股期间有除权的再做preEx操作
def profitCal(code, date, days):
    ExrightsToolObj = ExrightsTool()
    BasicToolObj = BasicTool()
    StockToolObj = StockTool()
    # 获得date后交易日期
    today = time.strftime('%Y-%m-%d %H:%M:%S',time.localtime(time.time()))
    flag, datelist = StockToolObj.get_trade_date(date, today)
    if len(datelist) < 1:
        return 0.0
    # 拿到除权因子
    flag, exData = ExrightsToolObj.getExrights(code)
    if flag is False:
        print 'get ex data error'
        return False
    # 拿到原始行情数据
    flag, oriItemList = BasicToolObj.getOriItemList(code, False, 'day', date+' 00:00:00', today+' 00:00:00')
    if flag is False:
        print 'get ori data error'
        return False
    oriItemList = oriItemList[1:days+1]
    if len(oriItemList) == 0:
        return 0.0
    buy_item = oriItemList[0]
    sell_item = oriItemList[-1]
    buy_price = buy_item[1]
    buy_date = buy_item[0][:10]
    sell_price = sell_item[2]
    sell_date = sell_item[0][:10]
    profit = (sell_price - buy_price) / buy_price
    # 除权判断
    for exdate in exData:
        if exdate > buy_date and exdate <= sell_date:
            # 开始除权
            exrights_data = exData[exdate]
            paixi = exrights_data["paixi"]
            peigu = exrights_data["peigu"]
            peigujia = exrights_data["peigujia"]
            songgu = exrights_data["songgu"]
            zenggu = exrights_data["zenggu"]
            buy_price = (buy_price - paixi + peigu * peigujia) / ( 1.0 + songgu + zenggu + peigu)
            profit = (sell_price - buy_price) / buy_price
            break
    return profit


class TrainSetProduce:
	def __init__(self, mem_mod=False, already_has_preData_inDesk=False, ignore_newstock_yiZiBan=True, filename='test.csv', codes=["300033", "000001"], label_name='closeProfit', period=['2013-01-01 00:00:00','2017-05-31 00:00:00'], shapeList = ['MACDShape'], shapeParam = {'MACDShape':{}}):
		self.shapeList = shapeList
		self.mem_mod = mem_mod
		self.filename = filename
		self.shapeParam = shapeParam
		self.label_name=label_name
		self.codes = codes
		self.period = period
		self.ignore_newstock_yiZiBan = ignore_newstock_yiZiBan
		self.codeFeatures = {}
		self.codeLabel = {}
		self.records = []
		self.StockPreExDataFile = "SPEDF(%s).data"%filename
		self.ExrightsToolObj = ExrightsTool()
		self.BasicToolObj = BasicTool()
		self.TecIndexRecogObj = TecIndexRecog()
		if already_has_preData_inDesk is False:
			self.saveStockData2Desk()


	# 保存计算好的pre复权数据到本地，加速计算
	# !!！附加新股开板前数据过滤！！！
	def saveStockData2Desk(self, ):
		print 'get ex right data ...'
		file_hd = open(self.StockPreExDataFile, 'w')
		i = 1
		for code in self.codes:
			i += 1
			if i%100 == 0:
				print i
			flag, itemList = self.ExrightsToolObj.getAdjItemList(code, "pre", "day", None, None)
			begin_index = self.NewStockKaiBan(code)
			itemList = itemList[begin_index:]
			if len(itemList) == 0:
				continue
			if flag is False:
				print 'get ex data error!! %s'%code
				return
			all_items = []
			for item in itemList:
				item_strs = [item[0]]
				for sub_item in item[1:]:
					item_strs.append(str(sub_item))
				item_str = ",".join(item_strs)
				all_items.append(item_str)
			one_code_data = code + ';' + ";".join(all_items)
			file_hd.write(one_code_data+'\n')
		file_hd.close()


	# 迭代得到所有前复权数据
	# file_SPEDF = open(self.StockPreExDataFile, 'r')
	# while True:
	#	code, itemList = self.iterStockData(file_SPEDF)
	#	if code == 'end':
	#		file_SPEDF.close()
	#       break
	# ！！！注意在迭代结束时关闭文件！！！
	def iterStockData(self, file_SPEDF):
		one_code_data = file_SPEDF.readline()
		if one_code_data == '':
			return 'end',None
		split_str = one_code_data.split(';')
		code = split_str[0]
		itemList = []
		for one_item in split_str[1:]:
			one_item_split = one_item.split(',')
			item = [one_item_split[0]]
			for one_item_split_one in one_item_split[1:]:
				item.append(float(one_item_split_one))
			itemList.append(item)
		return code, itemList
		

	# 在整个win下的收益(开盘买，收盘卖)，不考虑止盈损
	# 异常0，盈利1，亏损2
	# win_data: [item, item,.]
	def label_win_closeProfit(self, win_data):
		buy = win_data[0][1]
		sell = win_data[-1][2]
		if buy <= 0.0 or sell < 0.0:
			return 0.0
		profit = (sell - buy)/buy
		if profit <= 0.0:
			return profit
		else:
			return profit

	# 单个itemList的label计算通式
	# 返回 {datetime:label}，其中datetime指前一天的收盘买入时刻，而winN指买入后持有N天的label
	def label_cal(self, itemList, code, winN=3):
		length = len(itemList)
		if length < (winN+1):
			return {}
		# 拿到除权因子
		flag, exData = self.ExrightsToolObj.getExrights(code)
		if flag is False:
			print 'get ex data error'
			return False
		# 拿到原始行情数据
		flag, oriItemList = self.BasicToolObj.getOriItemList(code, False, 'day', None, None)
		if flag is False:
			print 'get ori data error'
			return False
		# 滑动窗口
		resultDict = {}
		#win_data = deque(itemList[0:winN])
		win_data = deque(oriItemList[0:winN])
		for i in range(winN,length):
			win_data.popleft()
			win_data.append(oriItemList[i])
			# 做决策的当天
			datetime = oriItemList[i-winN][0]
			# 持股期间
			date_s = oriItemList[i-winN+1][0][:10]
			date_e = oriItemList[i][0][:10]
			# 判断除权中断位置
			ex_info = None
			for ex_date in exData:
				# 有除权需要
				if ex_date > date_s and ex_date <= date_e:
					# 除权index
					for j in range(i-winN+1, i+1):
						date_ex = oriItemList[j][0][:10]
						if date_ex == ex_date:
							ex_info = [j, exData[ex_date]]
			# 计算局部除权数据
			if ex_info is not None:
				ex_index = ex_info[0]
				exrights_data = ex_info[1]
				paixi = exrights_data["paixi"]
				peigu = exrights_data["peigu"]
				peigujia = exrights_data["peigujia"]
				songgu = exrights_data["songgu"]
				zenggu = exrights_data["zenggu"]
				for j in range(i-winN+1, ex_index):
					cur_item = copy.deepcopy(oriItemList[j])
					# 前复权操作
					for k in range(1, 5):
						cur_item[k] = (cur_item[k] - paixi + peigu * peigujia) / ( 1.0 + songgu + zenggu + peigu)
					cur_item[5] = cur_item[5] * (1.0 + songgu + zenggu + peigu)
					win_data[j-i+winN-1] = cur_item
			paramDict = {'win_data': win_data}
			win_label = getattr(self, "label_win_" +self.label_name)(**paramDict)
			resultDict[datetime] = win_label
		return resultDict

	# 新股开板日期计算
	def NewStockKaiBan(self, code):
		if self.ignore_newstock_yiZiBan is False:
			return 0
		flag, itemList = self.BasicToolObj.getOriItemList(code, False, "day", None, None)
		if flag is False:
			print 'get ori stock data error'
			return 0
		lenItemList = len(itemList)
		if lenItemList <= 1:
			return len(itemList) - 1
		close_pre = itemList[0][2]
		for i in range(1, lenItemList):
			item = itemList[i]
			close = item[2]
			if (close - close_pre)/close_pre < 0.098:
				break
			close_pre = close
		# 至今未开板
		if lenItemList == (i+1):
			return lenItemList
		# 返回开板的第一天index(0开始) 
		return i

	# 计算各股票label
	# label_name的意义在于寻找label_win_closeProfit这样的函数
	def calLabel(self, winN=3):
		print 'label cal ...'
		if self.mem_mod is False:
			file_hd = open('calLabel.temp', 'w')
		i = 0
		file_SPEDF = open(self.StockPreExDataFile, 'r')
		while True:
			# 显示进度
			i += 1
			if i%100 == 0:
				print i
			code, itemList = self.iterStockData(file_SPEDF)
			if code == 'end':
				file_SPEDF.close()
				break
			self.codeLabel[code] = {}
			resultDict = self.label_cal(itemList, code, winN)
			# 过滤设定时间段以外的
			for datetime in resultDict:
				if datetime > self.period[0] and datetime < self.period[1]:
					self.codeLabel[code][datetime] = resultDict[datetime]
			# 省内存模式
			if self.mem_mod is False:
				for datetime in self.codeLabel[code]:
					label = self.codeLabel[code][datetime]
					line_str = code + ',' + datetime + ',' + str(label)
					file_hd.write(line_str + '\n')
				# 无数据，但是格式统一
				if len(self.codeLabel[code]) == 0:
					file_hd.write(code + ',' + 'no data\n')
				del self.codeLabel[code]
				
		
	# 计算各股票形态
	def calShape(self, ):
		print 'shape cal ...'
		if self.mem_mod is False:
			file_hd = open('calShape.temp', 'w')
		i = 0
		file_SPEDF = open(self.StockPreExDataFile, 'r')
		while True:
			# 显示进度
			i += 1
			if i%100 == 0:
				print i
			code, itemList = self.iterStockData(file_SPEDF)
			if code == 'end':
				file_SPEDF.close()
				break
			shape_datetimes = {}
			all_datetime = set([])
			for shapeName in self.shapeList:
				shape_datetimes[shapeName] = set([])
				shapeParam = self.shapeParam[shapeName]
				shapeParam['itemList'] = itemList
				# 返回的shapeResult 格式为： [ [datetime,f1,f2,],[], ]
				shapeResult = getattr(self.TecIndexRecogObj, shapeName)(**shapeParam)
				if not self.codeFeatures.has_key(code):
					self.codeFeatures[code] = {}
				for resultItem in shapeResult:
					datetime = resultItem[0]
					shape_datetimes[shapeName].add(datetime)
					all_datetime.add(datetime)
					# 过滤设定时间段以外的
					if datetime > self.period[0] and datetime < self.period[1]:
						if not self.codeFeatures[code].has_key(datetime):
							self.codeFeatures[code][datetime] = []
						# 末尾增量添加features
						self.codeFeatures[code][datetime].extend(resultItem[1:])
			# 求缺失datetime集合
			del_datetimes = set([])
			for shapeName in shape_datetimes:
				datetime_set = shape_datetimes[shapeName]
				del_datetimes = del_datetimes | (all_datetime - datetime_set)
			# 删除无公用datetime特征
			for del_datetime in del_datetimes:
				del self.codeFeatures[code][del_datetime]
			# 磁盘temp文件存储数据，节省内存
			if self.mem_mod is False:
				for datetime in self.codeFeatures[code]:
					features = self.codeFeatures[code][datetime]
					line_str = code + ',' + datetime + ',' + ','.join([str(k) for k in features])
					file_hd.write(line_str + '\n')
				# 没有数据，但是格式上要统一
				if len(self.codeFeatures[code]) == 0:
					file_hd.write(code + ',' + 'no data\n')
				del self.codeFeatures[code]
				

	# 将feature和label合并
	def CombineSet(self, ):
		print 'combine train data ...'
		# 低内存模式，合并同时写入文件
		if self.mem_mod is False:
			shape_file = open("calShape.temp", 'r')
			label_file = open("calLabel.temp", 'r')
			com_file = open(self.filename, 'w')
			com_file.write(self.shape_head_str() + '\n')
			start_shape_line = shape_file.readline()
			start_label_line = label_file.readline()
			i = 0
			while True:
				i += 1
				if i%100 == 0:
					print i
				# 同步读取形同code的数据
				cur_shape_line = start_shape_line
				cur_label_line = start_label_line
				start_shape_code = cur_shape_line[:6]
				start_label_code = cur_label_line[:6]
				#print start_shape_code, start_label_code
				# -----数据缺失验证
				start_shape_varify = cur_shape_line[7:-1]
				start_label_verify = cur_label_line[7:-1]
				#print start_shape_varify, start_label_verify
				if start_shape_varify == 'no data' and start_label_verify == 'no data':
					start_shape_line = shape_file.readline()
					start_label_line = label_file.readline()
					if start_shape_line == '' and start_label_line == '':
						break
					elif start_shape_line <> '' and start_label_line <> '':
						continue
					else:
						print 'not syn',(start_shape_line, start_label_line)
						break
				# ----------------
				if start_shape_code <> start_label_code:
					print start_shape_code,start_label_code,'code not equal!!!'
				shape_data = [cur_shape_line]
				label_data = [cur_label_line]
				while cur_shape_line[:6] == start_shape_code:
					shape_data.append(cur_shape_line)
					cur_shape_line = shape_file.readline()
				while cur_label_line[:6] == start_label_code:
					label_data.append(cur_label_line)
					cur_label_line = label_file.readline()
				# 处理code的数据，合并特征与类标
				shape_dict = {}
				for shape in shape_data:
					datetime = shape[7:26]
					# 跳过'，'与'\n'
					shape_dict[datetime] = shape[27:-1]
				label_dict = {}
				for label in label_data:
					datetime = label[7:26]
					# 跳过'，'与'\n'
					label_dict[datetime] = label[27:-1]
				for shape_datetime in shape_dict:
					label_str = label_dict.get(shape_datetime)
					if label_str is None:
						continue
					shape_str = shape_dict[shape_datetime]
					line_str = start_shape_code + ',' + shape_datetime +',' + shape_str + ',' + label_str
					# 将合并后的内容写入训练文件
					com_file.write(line_str + '\n')
				# 继续同步判断操作
				start_shape_line = cur_shape_line
				start_label_line = cur_label_line
				if start_shape_line == '' and start_label_line == '':
					break
				elif start_shape_line <> '' and start_label_line <> '':
					continue
				else:
					print 'not syn',(start_shape_line, start_label_line)
					break
			return
		self.records = []
		for code, Features in self.codeFeatures.items():
			Label = self.codeLabel[code]
			for datetime, features in Features.items():
				# 有feature，没有对应label的，过滤
				if Label.has_key(datetime):
					record = [code, datetime]
					for feature in features:
						record.append(feature)
					record.append(Label[datetime])
					self.records.append(record)

	def shape_head_str(self, ):
		# 组装csv文件头
		csv_head = ""
		comNameList = []
		for shapeName in self.shapeList:
			shapeParam = self.shapeParam[shapeName]
			if shapeParam.has_key('itemList'):
				del shapeParam['itemList']
			namelist = getattr(self.TecIndexRecogObj, shapeName+'Name')(**shapeParam)
			for oneName in namelist:
				#print shapeName, oneName
				comName = shapeName + '.' + oneName
				comNameList.append(comName)
		csv_head = ','.join(comNameList)
		csv_head = 'code,datetime,' + csv_head
		csv_head = csv_head + ',%s'%self.label_name
		return csv_head
		

	# record: ['000001', '2015-12-15 00:00:00', 1, 0, 1]
	def saveTrainSet(self, filename="trainSet_1.csv"):
		print 'save train data...'
		if self.mem_mod is False:
			return
		# 创建/删除文件
		file_hd = open(filename, 'w')
		csv_head = self.shape_head_str()
		file_hd.write(csv_head + '\n')
		# 写入record
		for record in self.records:
			record_str = record[0] + ',' + record[1]
			for record_item in record[2:]:
				record_str += ',' + str(record_item)
			file_hd.write(record_str+'\n')

	# 加载文件中的训练数据
	def loadTrainSet(self, ):
		print 'load train data ...'
		file_hd = open(self.filename, 'r')
		records = []
		# 跳过第一行
		for line in islice(file_hd, 1, None): 
			line_split = line.split(',')
			if len(line_split) <= 1:
				print 'record not complete: %s'%line
				continue
			line_split[-1] = line_split[-1].replace('\n','')
			records.append(line_split)
		return records


if __name__ == '__main__':
	StockToolObj = StockTool()
	#flag, codes = StockToolObj.get_all_stock_code('all')
	codes = ['600000']
	obj = TrainSetProduce(
			already_has_preData_inDesk=False,
			mem_mod=False,
			ignore_newstock_yiZiBan=True,
			filename="trainSet_allstock.csv",
			codes=codes,
			label_name='closeProfit', 
			period=['2007-01-01 00:00:00','2022-01-01 00:00:00'], 
			shapeList = ['MACDShape', 'MAShape'], 
			shapeParam = {'MACDShape':{}, 'MAShape':{}}
			)
	obj.calShape()
	#print obj.codeFeatures
	obj.calLabel(winN=3)
	#print obj.codeLabel
	obj.CombineSet()
	#print obj.records
	obj.saveTrainSet()
	#print obj.loadTrainSet()
